package prictice.hbase.state;

import org.apache.hadoop.hbase.client.Get;
import org.apache.hadoop.hbase.client.Put;
import org.apache.hadoop.hbase.client.Result;
import org.apache.hadoop.hbase.util.Bytes;
import org.apache.storm.trident.state.OpaqueValue;
import org.apache.storm.trident.state.Serializer;
import org.apache.storm.trident.state.StateFactory;
import org.apache.storm.trident.state.StateType;
import org.apache.storm.trident.state.map.IBackingMap;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

/**
 * @Package prictice.hbase.state
 * @Description:
 * @Author elwyn
 * @Date 2017/9/21 22:57
 * @Email elonyong@163.com
 */
@SuppressWarnings(value={"unchecked", "rawtypes"})
public class HbaseAggregateState<T> implements IBackingMap<T> {
    private HTableConnect hTableConnect;
    private Serializer<T> tSerializer;

    public HbaseAggregateState(TridentConfig tridentConfig) {
        this.tSerializer = tridentConfig.getSerializer();
        try {
            this.hTableConnect = new HTableConnect(tridentConfig);
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    public static StateFactory opaque(TridentConfig<OpaqueValue> config) {
        return new HBaseAggregateFactory(config, StateType.OPAQUE);
    }

    public static StateFactory transcational(TridentConfig<OpaqueValue> config) {
        return new HBaseAggregateFactory(config, StateType.TRANSACTIONAL);
    }

    public static StateFactory nontranscational(TridentConfig<OpaqueValue> config) {
        return new HBaseAggregateFactory(config, StateType.NON_TRANSACTIONAL);
    }

    @Override
    public List<T> multiGet(List<List<Object>> list) {
        List<Get> gets = new ArrayList<>(list.size());
        byte[] rk;
        byte[] cf;
        byte[] cq;
        for (List<Object> objects : list) {
            rk = Bytes.toBytes(String.valueOf(objects.get(0)));
            cf = Bytes.toBytes(String.valueOf(objects.get(1)));
            cq = Bytes.toBytes(String.valueOf(objects.get(2)));
            Get get = new Get(rk);
            gets.add(get.addColumn(cf, cq));
        }
        Result[] results = null;
        try {
            results = hTableConnect.gethTable().get(gets);
        } catch (IOException e) {
            e.printStackTrace();
        }
        List<T> rtn = new ArrayList<>(list.size());
        for (int i = 0; i < list.size(); i++) {
            cf = Bytes.toBytes(String.valueOf(list.get(i).get(1)));
            cq = Bytes.toBytes(String.valueOf(list.get(i).get(2)));
            Result result = results[i];
            if (result.isEmpty()) {
                rtn.add(null);
            } else {
                rtn.add((T) tSerializer.deserialize(result.getValue(cf, cq)));
            }
        }
        return rtn;
    }

    @Override
    public void multiPut(List<List<Object>> keys, List<T> values) {
        List<Put> puts = new ArrayList<>();
        byte[] rk;
        byte[] cf;
        byte[] cq;
        for (int i = 0; i < keys.size(); i++) {
            rk = Bytes.toBytes(String.valueOf(keys.get(i).get(0)));
            cf = Bytes.toBytes(String.valueOf(keys.get(i).get(1)));
            cq = Bytes.toBytes(String.valueOf(keys.get(i).get(2)));
            byte[] cv = tSerializer.serialize(values.get(i));
            Put put = new Put(rk);
            puts.add(put.add(cf, cq, cv));
        }

        try {
            hTableConnect.gethTable().put(puts);
            hTableConnect.gethTable().flushCommits();
        } catch (IOException e) {
            e.printStackTrace();
        }
    }
}
